#define _CRT_SECURE_NO_WARNINGS 1

/*
* Copyright (C) 2021  BryantDai
* All rights reserved.

* This program is free software: you can use, modify and/or redistribute it under the terms of the GNU General Public
* License as published by the Free Software Foundation.
*
* Notice : This file modify from vl_quickshift in VLFeat library
*
*
* Copyright (C) 2009-11 Andrea Vedaldi and Brian Fulkerson.
* All rights reserved.

* This file is part of the VLFeat library and is made available under the terms of the BSD license (see the COPYING file).
*/
/*********************************************************************************************************************

***********************************************************************************************************************/
/* Includes ----------------------------------------------------------------------------------------------------------*/
#include "QuickShift.h"

/* Private typedef ---------------------------------------------------------------------------------------------------*/

/* Private define ----------------------------------------------------------------------------------------------------*/
/* Private macro -----------------------------------------------------------------------------------------------------*/
/* Private variables -------------------------------------------------------------------------------------------------*/
/* Private function prototypes ---------------------------------------------------------------------------------------*/
/* Private functions -------------------------------------------------------------------------------------------------*/


/**********************************************************************************************************************
* @Name        : QuickShift
* @Brief       : Contructor funtion with parameters of QuickShift class
* @Param       : ratio      : Tradeoff between spatial consistency and color consistency.
*                kernelSize : The standard deviation of the parzen window density estimator.
*                maxDist    : The maximum distance between nodes in the quickshift tree. Used to cut links in the tree to form the segmentation.
* @Return      : None
* @Author      : Bryant
* @Create Date : Nov-11-2021
* @Mended by   :
* @Mend Date   :
* @Note        : (1)The superpixels extracted via quick-shift are controlled by three main parameters of ratio, KernelSize, and MaxDist.
*                   The ratio is a tradeoff between spatial consistency and color consistency. The KernelSize is the standard
*                   deviation of the parzen window density estimator, while MaxDist is the maximum distance between nodes in the
*                   quick-shift tree, which used to cut links in the tree to form the segmentation.
***********************************************************************************************************************/
QuickShift::QuickShift(double ratio, double kernelSize, double maxDist)
{
    this->m_ratio = ratio;
    this->m_maxDist = maxDist;
    this->m_kernelSize = kernelSize;
}


/**********************************************************************************************************************
* @Name        : QuickShiftSegment
* @Brief       : Produce a quickshift segmentation of a grayscale or color image
* @Param       : srcImg : Input image,may be RGB or Grayscale. RGB images are first converted to LAB.
*                pntIdxArr : Each element's (pixel) index to the parent element in the forest
*                distMat   : Contains the corresponding branch length
* @Return      : segImg : A color image where each pixel is labeled by the mean color in its region.
* @Author      : Bryant
* @Create Date : Nov-12-2021
* @Mended by   :
* @Mend Date   :
* @Note        : (1) Add less than one pixel noise to break ties caused by constant regions in an arbitrary fashon.
*                    (Modify from vl_quickshift in VLFeat library )
*                (2) Small ratio gives more importance to the spatial component. Note that distance calculations happen
*                    in unnormalized image coordinates, so RATIO should be adjusted to compensate for larger images.
***********************************************************************************************************************/
cv::Mat QuickShift::QuickShiftSegment(cv::Mat& srcImg, vector2L& pntIdxArr, vector2D& distMat)
{
    cv::Mat doubleImg;
    srcImg.convertTo(doubleImg, CV_32FC3, 1.0 / 255.0);       /* Convert image to double precision                      */

  /*  cv::Mat randMat(srcImg.rows, srcImg.cols, CV_32FC3);
    cv::randn(randMat, 0, 1);
    doubleImg += randMat / 255;  */                           /* See Note (1)                                           */
    cv::cvtColor(doubleImg, doubleImg, cv::COLOR_BGR2Lab);    /* From RGB to Lab feature space                          */

    doubleImg = doubleImg * this->m_ratio;                    /* See Note(2)                                            */
    QuickShiftProcess(doubleImg, m_kernelSize, m_maxDist, distMat, pntIdxArr);
    ImgSegment(doubleImg, pntIdxArr, distMat);

    cv::cvtColor(doubleImg, doubleImg, cv::COLOR_Lab2BGR);    /* Convert image back to RGB                              */
    doubleImg.convertTo(doubleImg, CV_8UC3, 255.0 / 1.0);
    return doubleImg;
}


/**********************************************************************************************************************
* @Name        : ImgSegment
* @Brief       : Segment the image basic on the labels
* @Param       : srcImg      : Input image,may be RGB or Grayscale.
*                pntIdxArray : Store the best neighbor of the corresponding pixel point
*                distMat     : Contains the corresponding branch length
* @Return      : None
* @Author      : Bryant
* @Create Date : Nov-14-2021
* @Mended by   :
* @Mend Date   :
* @Note        :
***********************************************************************************************************************/
void QuickShift::ImgSegment(cv::Mat& srcImg, vector2L& pntIdxArray, vector2D& distMat)
{
    int clusterNum = 0;
    int rowMax = (int)pntIdxArray.size();
    int colMax = (int)pntIdxArray[0].size();
    std::vector<uint32_t> medoid;

    /* Get the labels of all the pixels in the image                                                                    */
    for (int rowIndex = 0; rowIndex < rowMax; rowIndex++)
    {
        for (int colIndex = 0; colIndex < colMax; colIndex++)
        {
            /* Follow the parents list to the root nodes (where nothing changes)                                            */
            uint32_t med = pntIdxArray[rowIndex][colIndex];
            if (QUICKSHIFT_INF == distMat[rowIndex][colIndex])     /* Those are the tree roots                              */
            {
                clusterNum++;
                medoid.push_back(rowIndex * colMax + colIndex);
            }
            else
            {
                while (med != pntIdxArray[med / colMax][med % colMax])  /* Find the tree root this pixel attach to            */
                {
                    med = pntIdxArray[med / colMax][med % colMax];
                }
            }

            pntIdxArray[rowIndex][colIndex] = med;
        } /* End of for (int colIndex = 0; colIndex < srcImg.cols; colIndex++)                                            */
    } /* End of for (int rowIndex = 0; rowIndex < srcImg.rows; rowIndex++)                                              */

    LabelImage(srcImg, pntIdxArray, medoid);
}


/**********************************************************************************************************************
* @Name        : LabelImage
* @Brief       : Label image basic on the result getting from quick shift segmentation.Labels 'srcImg' with the average
*                color from 'srcImg' of  each cluster indicated by labels.
* @Param       : srcImg : Input image
*                labels : Cluster labels
*                segNum : Number of section
* @Return      : None
* @Author      : Bryant
* @Create Date : Nov-14-2021
* @Mended by   :
* @Mend Date   :
* @Note        :
***********************************************************************************************************************/
void QuickShift::LabelImage(cv::Mat& srcImg, vector2L& labels, std::vector<uint32_t>& medoid)
{
    int segNum = (int)medoid.size();
    std::vector<int> attachedNum(segNum, 0);
    std::vector<cv::Vec3f> attachedValue(segNum, 0);
    int rowMax = (int)labels.size();
    int colMax = (int)labels[0].size();

    for (int rowIndex = 0; rowIndex < rowMax; rowIndex++)
    {
        for (int colIndex = 0; colIndex < colMax; colIndex++)
        {
            int medIndex = (int)(std::find(medoid.begin(), medoid.end(), labels[rowIndex][colIndex]) - medoid.begin());
            attachedNum[medIndex]++;                  /* Record the pixel number attaching to each tree root                */
            attachedValue[medIndex] += srcImg.at<cv::Vec3f>(rowIndex, colIndex); /* Sum all pixel value int the same tree up*/
        }
    }

#pragma omp parallel for
    /* Average color from 'srcImg' of each cluster indicated by labels.                                                 */
    for (int size = 0; size < attachedNum.size(); ++size)
    {
        attachedValue[size] /= attachedNum[size];
    }

#pragma omp parallel for
    for (int rowIndex = 0; rowIndex < rowMax; rowIndex++)   /* Labels 'srcImg' with the average color of each cluster   */
    {
        for (int colIndex = 0; colIndex < colMax; colIndex++)
        {
            int medIndex = (int)(std::find(medoid.begin(), medoid.end(), labels[rowIndex][colIndex]) - medoid.begin());
            srcImg.at<cv::Vec3f>(rowIndex, colIndex) = attachedValue[medIndex];
        }
    }
}


/**********************************************************************************************************************
* @Name        : QuickShiftProcess
* @Brief       : Perform quickshift to obtain the segmentation tree
* @Param       : srcImg     : Input image
*                kernelSize : The bandwidth of the Parzen window estimator of the density
*                maxDist    : Controls the maximum L2 distance between neighbors that should be linked
*                gapLenMat  : Contains the corresponding branch length
*                pntIdxArr  : Store the best neighbor of the corresponding pixel point
* @Return      : None
* @Author      : Bryant
* @Create Date : Nov-12-2021
* @Mended by   :
* @Mend Date   :
* @Note        :
***********************************************************************************************************************/
void QuickShift::QuickShiftProcess(cv::Mat& srcImg, double kernelSize, double maxDist, vector2D& gapLenMat, vector2L& pntIdxArr)
{
    vector2D weightMatrix = BuildWeightMatrix(srcImg, kernelSize);
    FindBestNeighbors(srcImg, weightMatrix, pntIdxArr, gapLenMat, maxDist);
}


/**********************************************************************************************************************
* @Name        : FindBestNeighbors
* @Brief       : Quickshift assigns each i to the closest j which has an increase in the density
* @Param       : srcImg    : Input image
*                weightMat : The weight of each pair of points
*                maxDist   : Controls the maximum L2 distance between neighbors that should be linked
*                parentMat : Store the best neighbor of the corresponding pixel point
* @Return      : None
* @Author      : Bryant
* @Create Date : Nov-12-2021
* @Mended by   :
* @Mend Date   :
* @Note        : (1) Quickshift assigns each i to the closest j which has an increase in the density (E).
*                    If there is no j s.t. Ej > Ei, then dists_i == inf (a root node in one of the trees of merges).
***********************************************************************************************************************/
void QuickShift::FindBestNeighbors(cv::Mat& srcImg, vector2D& weightMat, vector2L& parentMat, vector2D& distMat, double maxDist)
{
    int maxDistance = (int)std::ceil(maxDist);
    double maxDistPow2 = maxDist * maxDist;

#pragma omp parallel for
    for (int rowIndex = 0; rowIndex < srcImg.rows; rowIndex++)
    {
        for (int colIndex = 0; colIndex < srcImg.cols; colIndex++)
        {
            double E0 = weightMat[rowIndex][colIndex];
            double dBest = QUICKSHIFT_INF;
            int rowBest = rowIndex;
            int colBest = colIndex;

            /* The valid distance range in spatial domain                                                                   */
            int rowStart = std::max(0, rowIndex - maxDistance);
            int colStart = std::max(0, colIndex - maxDistance);
            int rowEnd = std::min(srcImg.rows - 1, rowIndex + maxDistance);
            int colEnd = std::min(srcImg.cols - 1, colIndex + maxDistance);
            for (int validRowIdx = rowStart; validRowIdx < rowEnd; ++validRowIdx)
            {
                for (int validColIdx = colStart; validColIdx < colEnd; ++validColIdx)
                {
                    if (weightMat[validRowIdx][validColIdx] > E0)
                    {
                        double D_ij = CalPointsDist(srcImg, rowIndex, validRowIdx, colIndex, validColIdx);
                        if ((D_ij < maxDistPow2) && (D_ij < dBest))   /* See Note (1)                                             */
                        {
                            dBest = D_ij;
                            rowBest = validRowIdx;
                            colBest = validColIdx;
                        }
                    }
                }
            } /* End of for (int validRowIdx = rowStart; validRowIdx < rowEnd; ++validRowIdx)                               */

            parentMat[rowIndex][colIndex] = srcImg.cols * rowBest + colBest; /* Parents is the index of the best pair       */
            if (dBest != QUICKSHIFT_INF)
            {
                distMat[rowIndex][colIndex] = std::sqrt(dBest);
            }
            else
            {
                distMat[rowIndex][colIndex] = QUICKSHIFT_INF;
            }
        } /* End of for (int colIndex = 0; colIndex < srcImg.cols; colIndex++)                                            */
    } /* End of for (int rowIndex = 0; rowIndex < srcImg.rows; rowIndex++)                                              */
}


/**********************************************************************************************************************
* @Name        : BuildWeightMatrix
* @Brief       : Compute the weight of each pair of points, according to the window chosen
* @Param       : srcImg : Input image
*                sigma  : The standard deviation of the kernel used in the Parzen density estimate.
* @Return      : None
* @Author      : Bryant
* @Create Date : Nov-11-2021
* @Mended by   :
* @Mend Date   :
* @Note        : (1) The Parzen window used in this implementation is k(x) = exp(-D_ij/2(sigma^2))
***********************************************************************************************************************/
vector2D QuickShift::BuildWeightMatrix(cv::Mat& srcImg, double sigma)
{
    int spatialRadius = (int)std::ceil(3 * sigma);
    vector2D weigthMatrix(srcImg.rows, std::vector<double>(srcImg.cols, 0));

#pragma omp parallel for
    for (int rowIndex = 0; rowIndex < srcImg.rows; rowIndex++)
    {
        for (int colIndex = 0; colIndex < srcImg.cols; colIndex++)
        {
            /* The valid distance range in spatial domain                                                                   */
            int rowStart = std::max(0, rowIndex - spatialRadius);
            int colStart = std::max(0, colIndex - spatialRadius);
            int rowEnd = std::min(srcImg.rows - 1, rowIndex + spatialRadius);
            int colEnd = std::min(srcImg.cols - 1, colIndex + spatialRadius);

            /* For each pixel in the window compute the distance between it and the source pixel                            */
            for (int validRowIdx = rowStart; validRowIdx < rowEnd; ++validRowIdx)
            {
                for (int validColIdx = colStart; validColIdx < colEnd; ++validColIdx)
                {
                    double D_ij = CalPointsDist(srcImg, rowIndex, validRowIdx, colIndex, validColIdx);
                    double F_ij = std::exp(-D_ij / (2 * sigma * sigma));   /* See Note (1)                                      */
                    weigthMatrix[rowIndex][colIndex] += F_ij;              /* Make distance a similarity                        */
                }
            } /* End of for (int rowI = rowStart; rowI < rowEnd; ++rowI)                                                    */
        } /* End of for (int colIndex = 0; colIndex < srcImg.cols; colIndex++)                                            */
    } /* End of for (int rowIndex = 0; rowIndex < srcImg.rows; rowIndex++)                                              */

    return weigthMatrix;
}


/**********************************************************************************************************************
* @Name        : CalPointsDist
* @Brief       : Computes The L2 distance between two points
* @Param       : srcImg     : Input image
*                pts1RowIdx : The row index of point1
*                pts2RowIdx : The row index of point2
*                pts1ColIdx : The column index of point1
*                pts2ColIdx : The column index of point2
* @Return      : The distance as described above
* @Author      : Bryant
* @Create Date : Nov-11-2021
* @Mended by   :
* @Mend Date   :
* @Note        : (1) The quick shift is build in the joint spatial and color dimensions.
***********************************************************************************************************************/
double QuickShift::CalPointsDist(cv::Mat& srcImg, int pts1RowIdx, int pts2RowIdx, int pts1ColIdx, int pts2ColIdx)
{
    double distance = 0;

    int rowDist = pts2RowIdx - pts1RowIdx;
    int colDist = pts2ColIdx - pts1ColIdx;
    distance += rowDist * rowDist + colDist * colDist;

    distance += std::pow(cv::norm(srcImg.at<cv::Vec3f>(pts1RowIdx, pts1ColIdx), srcImg.at<cv::Vec3f>(pts2RowIdx, pts2ColIdx), cv::NORM_L2), 2);

    return distance;
}


/************************************************* End of file ********************************************************/

